from torch import nn
import torch


class TimeSeriesClassifier(nn.Module):
    def __init__(self, n_features, hidden_dim=64, output_size=1):
        super().__init__()
        self.lstm = nn.LSTM(input_size=n_features, hidden_size=hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_size)  # output_size classes

    def forward(self, x):
        x, _ = self.lstm(x)  # LSTM层
        x = x[:, -1, :]  # 只取LSTM输出中的最后一个时间步
        x = self.fc(x)  # 通过一个全连接层
        return x

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# 下面函数接收的参数中，input为一个numpy可识别的数组，后面的参数暂时都为默认参数
def get_output(input, n_features = 14,output_size = 4,param ='LSTM_best_model.pth'):
    model = TimeSeriesClassifier(n_features=n_features, output_size=output_size)
    checkpoint = torch.load(param, weights_only=False, map_location=device)
    #model = torch.load('LSTM_best_model.pth')
    model.load_state_dict(checkpoint['state_dict']) 
    model.eval()
    model.to(device)
    input = input.reshape(1, 50, 14)
    data = torch.tensor(input).float().to(device)
    outputs = model(data)
    _, predicted = torch.max(outputs, 1)
    results = predicted.cpu()
    output = (results.numpy()).tolist()

    return output
